'''
    This is used to store the correspondence between the basic data set and the feature (using h 5)
'''

import os

import h5py
import torch

from src.CNNModel.Net.VGG16 import SE_VGG
from src.CNNModel.Net.ResNet import Net
import matplotlib.image as mpimg
import matplotlib.pyplot as plt
import LoadImage as lim
import numpy as np
import torchvision.models as models

# Get the path of all pictures
def get_imlist(path):
    return [os.path.join(path, f) for f in os.listdir(path) if f.endswith('.jpg')]

if __name__ == '__main__':

    features = [] # 存储特征值
    paths = [] # 存储路径

    datapath = 'DB/pictures' # 数据地址
    output = 'FeaturesAndPaths/vgg_featureCNN.h5'
    img_list = get_imlist(datapath) # 获得图像数据
    print(img_list)
    print('------------------------------------------\n'
          '       特征提取开始          \n'
          '------------------------------------------')

    # model = SE_VGG(50) # (*,3,224,224)
    # model = Net()
    model = models.vgg16(pretrained=False)
    pthfile = 'Model/state/vgg16-397923af.pth' #预训练模型的地址
    model.load_state_dict(torch.load(pthfile))#将预训练的参数权重加载到新的模型之中
    model.eval() # 不启用 BatchNormalization 和 Dropout
    print(model)

    for idx,image_path in enumerate(img_list):
        '''
        # 测试
        print(image_path)
        queryImg = mpimg.imread(image_path)
        plt.title("Query Image")
        plt.imshow(queryImg)
        plt.show()
        '''
        # 顺序处理
        img = lim.MyLoader(image_path) # 加载图片
        img = lim.transform(img) # 图像增强处理
        img = img.view(1,3,224,224) # 到四维转换
        # print(img)
        norm_feat = model(img) # 得到特征，就可以改变网络
        print(norm_feat.size(1))
        norm_feat = norm_feat.view(norm_feat.size(1))
        # print(norm_feat)
        features.append(norm_feat.detach().numpy()) # 存储特征
        paths.append(image_path) # 存储相应的路径
        print(f'-------- Path : {image_path} 处理完毕-------')

    features = np.array(features)
    print(features,'\n',paths)

    print('------------------------------------------\n'
          '  编写特征提取结果   \n'
          '------------------------------------------')

    # print(np.string_(paths))
    h5f = h5py.File(output,'w') # create h5 file
    h5f.create_dataset('features',data=features)
    h5f.create_dataset('paths',data=np.string_(paths))
    h5f.close()

    print('------------------------------------------\n'
          '           写入成功           \n'
          '------------------------------------------')



